import os
import shutil

import numpy as np
import scipy.io.wavfile as wavfile
import tensorflow as tf

# ==================== config for spectrogram ==========
os.environ["CUDA_VISIBLE_DEVICES"] = '2'

wav_dir = '/home/ddy/projects/emotions/iemocap_4emo_wav'
spectrogram_dir = '/home/ddy/projects/emotions/iemocap_4emo_spectrogram'

# wav_dir = '/Users/d/Project/emotions/data/iemocap_4emo_wav_simple'
# spectrogram_dir = '/Users/d/Project/emotions/data/iemocap_4emo_spectrogram_simple'

count = 0


# =================== convert wav to spectrogram =====================
def cal_spectrogram(signals, frame_size, step_frac):
    """signals: tf.placeholder [batch_size, signal_length]"""
    # [batch_size, signal_length]
    # signals = tf.placeholder(tf.float32, shape=[None, None], name='signals')
    frame_step = int(frame_size * step_frac)
    stfts = tf.contrib.signal.stft(signals, frame_length=frame_size,
                                   frame_step=frame_step,
                                   fft_length=1600,
                                   window_fn=tf.contrib.signal.hamming_window,
                                   pad_end=True)
    magnitude_spectrogram = tf.abs(stfts)
    log_offset = 1e-6
    log_spectrogram = 10 * tf.log(magnitude_spectrogram + log_offset)
    return log_spectrogram


def cal_save_session(session_name, frame_time=0.04, step_frac=0.25):
    global count
    signals = tf.placeholder(tf.float32, [1, None])
    sentence_label_filepath = os.path.join(wav_dir,
                                           session_name + '_sentence_label')
    spectr_session_path = os.path.join(spectrogram_dir, session_name)
    if os.path.exists(spectr_session_path):
        shutil.rmtree(spectr_session_path)
    os.makedirs(spectr_session_path)
    # wav_paths = list()
    with tf.Session() as sess:
        with open(sentence_label_filepath, 'r') as s_f:
            for line in s_f:
                eles = line.split()
                if len(eles) == 2:
                    wav_path = os.path.join(wav_dir, session_name,
                                            eles[0] + '.wav')
                    rate, data = wavfile.read(wav_path)
                    data = data.reshape([1, -1])
                    frame_size = int(rate * frame_time)
                    l_spectr = cal_spectrogram(signals, frame_size, step_frac)
                    spectrogram = \
                        l_spectr.eval(feed_dict={signals: data}, session=sess)[
                            0]
                    time_scale, freq_scale = spectrogram.shape
                    useful_spectrogram = spectrogram[:, :int(freq_scale / 2)]
                    count += 1
                    print('count:{0}'.format(count))
                    print(eles[0], time_scale, freq_scale)
                    spectr_path = os.path.join(spectrogram_dir, session_name,
                                               eles[0] + '.npy')
                    np.save(spectr_path, useful_spectrogram)


def cal_save_sessions():
    if os.path.exists(spectrogram_dir):
        shutil.rmtree(spectrogram_dir)
    os.makedirs(spectrogram_dir)
    session_names = ['Session1', 'Session2', 'Session3', 'Session4', 'Session5']

    for session_name in session_names:
        cal_save_session(session_name)


def main():
    cal_save_sessions()


if __name__ == '__main__':
    main()
